package com.xlh.dokka.svc.service;

import com.google.common.collect.HashBasedTable;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.google.common.collect.Table;
import com.xlh.dokka.api.DokkaConfigService;
import com.xlh.dokka.api.DokkaContainerService;
import com.xlh.dokka.api.bo.DokkaClusterInfoBO;
import com.xlh.dokka.api.bo.DokkaContainerBO;
import com.xlh.dokka.api.constant.DokkaContainerConstants;
import com.xlh.dokka.api.dto.DokkaContainerCreateDTO;
import com.xlh.dokka.api.exception.DokkaException;
import com.xlh.dokka.api.exception.DokkaStartException;
import com.xlh.dokka.svc.core.client.DockerClient;
import com.xlh.dokka.svc.core.client.Result;
import com.xlh.dokka.svc.core.model.*;
import com.xlh.dokka.svc.core.novnc.LocalNovncConnector;
import com.xlh.dokka.svc.core.novnc.NovncConnector;
import com.xlh.dokka.svc.core.novnc.RemoteNovncConnector;
import com.xlh.dokka.svc.core.scheduler.Scheduler;
import com.xlh.dokka.svc.core.scheduler.filter.Filter;
import com.xlh.dokka.svc.core.scheduler.filter.HealthFilter;
import com.xlh.dokka.svc.core.scheduler.strategy.ResourcesSpreadStrategy;
import com.xlh.dokka.svc.core.scheduler.strategy.Strategy;
import com.xlh.dokka.svc.dao.DokkaContainerMapper;
import com.xlh.dokka.svc.properties.DockerProperties;
import com.xlh.dokka.svc.properties.NovncProperties;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationListener;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * @author cheer
 */
@Service
@Slf4j
public class DokkaContainerServiceImpl implements DokkaContainerService, ApplicationListener<ApplicationReadyEvent> {

    private Map<String, Node> nodes;

    private Table<String, Integer, GpuNode> gpuNodes;

    private Scheduler scheduler;

    private NovncConnector connector;

    /**
     * 控制同步节点使用情况的读写锁，保证同步操作与创建、启动、停止等操作是互斥的
     * 创建、启动、停止容器时获取读锁，同步节点使用情况时需获取写锁
     */
    private ReadWriteLock syncLock = new ReentrantReadWriteLock();

    /**
     * 控制创建容器选择节点的读写锁，保证选择节点操作时，节点使用情况是不变的
     * 修改资源使用情况时获取读锁，选择节点时需获取写锁
     */
    private ReadWriteLock selectLock = new ReentrantReadWriteLock();

    @Autowired
    private ThreadPoolTaskScheduler taskScheduler;

    @Autowired
    private ThreadPoolTaskExecutor taskExecutor;

    @Autowired
    private DokkaConfigService dokkaConfigService;

    @Autowired
    private DokkaContainerMapper dokkaContainerMapper;

    @Autowired
    private DockerProperties dockerProperties;

    @Autowired
    private NovncProperties novncProperties;

    @Override
    public void onApplicationEvent(ApplicationReadyEvent event) {
        initializeNodes();
        initializeNodeUsage();
        initializeGpuUsageCount();
        initializeScheduler();
        initializeNovncConnector();
        initializeDeleteErrorContainer();
    }

    List<Node> getNodes() {
        return new ArrayList<>(nodes.values());
    }

    @Override
    public DokkaClusterInfoBO getClusterInfo() {
        int totalCpus = 0;
        double usedCpus = 0;
        long totalMemory = 0;
        long usedMemory = 0;
        long totalDisk = 0;
        long usedDisk = 0;

        for (Node node : nodes.values()) {
            ResourceUsage resourceUsage = node.getResourceUsage();
            if (resourceUsage != null) {
                totalCpus += resourceUsage.getTotalCpus();
                usedCpus += resourceUsage.getCpuUsage() * resourceUsage.getTotalCpus();
                totalMemory += resourceUsage.getTotalMemory();
                usedMemory += resourceUsage.getUsedMemory();
                totalDisk += resourceUsage.getTotalDisk();
                usedDisk += resourceUsage.getUsedDisk();
            }
        }

        return DokkaClusterInfoBO.builder()
                .totalCpus(totalCpus)
                .usedCpus(Math.round(usedCpus * 100) / 100.0)
                .totalMemory(totalMemory)
                .usedMemory(usedMemory)
                .totalDisk(totalDisk)
                .usedDisk(usedDisk)
                .build();
    }

    @Override
    public DokkaContainerBO createContainer(DokkaContainerCreateDTO dokkaContainerCreateDTO) throws DokkaException {
        Container preContainer = dokkaContainerMapper.getById(dokkaContainerCreateDTO.getDokkaId());

        ContainerConfig containerConfig = ContainerConfig.builder()
                .image(dokkaContainerCreateDTO.getImage())
                .cpu((double) dokkaContainerCreateDTO.getCpu())
                .memory((long) dokkaContainerCreateDTO.getMemory() << 30)
                .hostname(dokkaContainerCreateDTO.getHostname())
                .mounts(dokkaContainerCreateDTO.getMounts())
                .exposedPorts(dokkaContainerCreateDTO.getExposedPorts())
                .host(preContainer != null ? preContainer.getHost() : null)
                .build();

        Container container;
        syncLock.readLock().lock();
        try {
            if (dokkaContainerCreateDTO.getWithGpu()) {
                container = createContainerWithGpu(containerConfig, dokkaContainerCreateDTO.getGpuPercent());
            } else {
                container = createContainer(containerConfig, true);
            }
        } finally {
            syncLock.readLock().unlock();
        }

        return DokkaContainerBO.builder()
                .dokkaId(container.getId())
                .name(container.getName())
                .ip(container.getIp())
                .status(DokkaContainerConstants.STATUS_RUNNING)
                .cpu(dokkaContainerCreateDTO.getCpu())
                .memory(dokkaContainerCreateDTO.getMemory())
                .novnc(container.getNovnc())
                .build();
    }

    @Override
    public void deleteContainer(Long dokkaId) {
        Container container = getContainerById(dokkaId);
        if (container == null) {
            return;
        }

        if ("running".equals(container.getStatus())) {
            syncLock.readLock().lock();
            try {
                decrNodeUsage(container);
                if (DockerClient.stopContainer(container.getHost(), container.getDockerId()) != Result.NO_ERROR) {
                    incrNodeUsage(container);
                }
            } finally {
                syncLock.readLock().unlock();
            }
        }

        if ("unknown".equals(container.getStatus()) ||
                !DockerClient.deleteContainer(container.getHost(), container.getDockerId())) {
            retryDeleteContainer(container.getHost(), container.getDockerId());
        }

        try {
            dokkaContainerMapper.deleteById(dokkaId);
        } catch (Exception e) {
            log.error("delete container {} failed", dokkaId, e);
        }

        connector.deleteToken(dokkaId);
    }

    @Override
    public Boolean updateContainer(Long dokkaId, Integer cpu, Integer memory) {
        Container container = getContainerById(dokkaId);
        if (container == null || "unknown".equals(container.getStatus())) {
            return false;
        }

        if ("running".equals(container.getStatus())) {
            syncLock.readLock().lock();
            try {
                decrNodeUsage(container);
                Result result = DockerClient.stopContainer(container.getHost(), container.getDockerId());
                if (result != Result.NO_ERROR) {
                    incrNodeUsage(container);
                }
                if (result == Result.SERVER_ERROR || result == Result.NO_SUCH_CONTAINER) {
                    return false;
                }
            } finally {
                syncLock.readLock().unlock();
            }
        }

        return DockerClient.updateContainer(container.getHost(), container.getDockerId(),
                (double) cpu, (long) memory << 30);
    }

    @Override
    public DokkaContainerBO getContainer(Long dokkaId) {
        Container container = getContainerById(dokkaId);
        if (container == null || "unknown".equals(container.getStatus())) {
            return null;
        }

        if ("running".equals(container.getStatus())) {
            if (container.getVncExport() == null) {
                log.error("container with dokkaId {} vnc export failed", dokkaId);
                return null;
            }
            container.setNovnc(connector.getNovnc(container.getId(), container.getHost() + ":" + container.getVncExport()));
        }

        return DokkaContainerBO.builder()
                .dokkaId(container.getId())
                .name(container.getName())
                .ip(container.getIp())
                .status("running".equals(container.getStatus()) ? DokkaContainerConstants.STATUS_RUNNING : DokkaContainerConstants.STATUS_EXITED)
                .cpu(container.getCpu().intValue())
                .memory((int) (container.getMemory() >> 30))
                .novnc(container.getNovnc())
                .build();
    }

    @Override
    public DokkaContainerBO startContainer(Long dokkaId) throws DokkaException {
        Container container = getContainerById(dokkaId);
        if (container == null || "unknown".equals(container.getStatus())) {
            throw new DokkaStartException("创建实验环境异常");
        }

        if (!"running".equals(container.getStatus())) {
            syncLock.readLock().lock();
            try {
                Node node;
                selectLock.writeLock().lock();
                try {
                    Integer containerQuota = dokkaConfigService.getContainerQuota();
                    if (containerQuota != null && getRunningCount() >= containerQuota) {
                        throw new DokkaException("运行容器数已达到设置最大数，无法创建更多实验环境");
                    }

                    node = nodes.get(container.getHost());
                    if (node != null && node.getNodeUsage() != null) {
                        node.getNodeUsage().update(container.getCpu(), container.getMemory(), 1);
                    }
                    if (container.getGpu() >= 0) {
                        GpuNode gpuNode = gpuNodes.get(container.getHost(), container.getGpu());
                        if (gpuNode != null) {
                            gpuNode.getUsageCount().getAndIncrement();
                        }
                    }
                } finally {
                    selectLock.writeLock().unlock();
                }

                if (node != null && !node.tryAcquire()) {
                    decrNodeUsage(container);
                    throw new DokkaException("正在启动实验环境的人数过多，请稍后在进入实验环境");
                }

                try {
                    Result result = DockerClient.startContainer(container.getHost(), container.getDockerId());
                    if (result != Result.NO_ERROR) {
                        decrNodeUsage(container);
                    }
                    if (result == Result.SERVER_ERROR || result == Result.NO_SUCH_CONTAINER || result == Result.FORBIDDEN) {
                        throw new DokkaStartException("创建实验环境异常");
                    }
                } finally {
                    if (node != null) {
                        node.releaseAcquire();
                    }
                }
            } finally {
                syncLock.readLock().unlock();
            }
        }

        Container inspectContainer = DockerClient.inspectContainer(container.getHost(), container.getDockerId());
        if (inspectContainer == null || !"running".equals(inspectContainer.getStatus())) {
            log.error("inspect container failed");
            throw new DokkaStartException("创建实验环境异常");
        }

        String novnc = connector.getNovnc(container.getId(), container.getHost() + ":" + inspectContainer.getVncExport());
        if (novnc == null) {
            log.error("get novnc failed");
            throw new DokkaException("创建实验环境异常");
        }

        return DokkaContainerBO.builder()
                .dokkaId(container.getId())
                .name(inspectContainer.getName())
                .ip(inspectContainer.getIp())
                .status(DokkaContainerConstants.STATUS_RUNNING)
                .cpu(inspectContainer.getCpu().intValue())
                .memory((int) (inspectContainer.getMemory() >> 30))
                .novnc(novnc)
                .build();
    }

    @Override
    public Boolean stopContainer(Long dokkaId) {
        Container container = getContainerById(dokkaId);
        if (container == null || "unknown".equals(container.getStatus())) {
            return false;
        }
        if ("exited".equals(container.getStatus())) {
            return true;
        }

        syncLock.readLock().lock();
        try {
            decrNodeUsage(container);
            Result result = DockerClient.stopContainer(container.getHost(), container.getDockerId());
            if (result != Result.NO_ERROR) {
                incrNodeUsage(container);
            }
            return result == Result.NO_ERROR || result == Result.CONTAINER_ALREADY_DONE;
        } finally {
            syncLock.readLock().unlock();
        }
    }

    @Override
    public Boolean execContainer(Long dokkaId, String command) {
        Container container = dokkaContainerMapper.getById(dokkaId);
        if (container == null) {
            return false;
        }

        String execId = DockerClient.createExec(container.getHost(), container.getDockerId(), command);
        return execId != null && DockerClient.startExec(container.getHost(), execId);
    }

    @Override
    public Boolean isSameHost(Long preDokkaId, Long afterDokkaId) {
        Container preContainer = dokkaContainerMapper.getById(preDokkaId);
        Container afterContainer = dokkaContainerMapper.getById(afterDokkaId);

        if (preContainer == null || afterContainer == null) {
            return false;
        }

        return Objects.equals(preContainer.getHost(), afterContainer.getHost());
    }

    private void initializeNodes() {
        Set<String> hosts = new HashSet<>();
        if (CollectionUtils.isNotEmpty(dockerProperties.getHosts())) {
            hosts.addAll(dockerProperties.getHosts());
        }

        Set<String> gpuHosts = new HashSet<>();
        if (CollectionUtils.isNotEmpty(dockerProperties.getGpuHosts())) {
            gpuHosts.addAll(dockerProperties.getGpuHosts());
        }

        nodes = new HashMap<>();
        for (String ip : Sets.union(hosts, gpuHosts)) {
            Node node = new Node(ip);
            node.initialize();
            nodes.put(ip, node);
        }

        gpuNodes = HashBasedTable.create();
        for (String ip : gpuHosts) {
            Node node = nodes.get(ip);
            node.setGpu(true);
            for (int serialNumber = 0; serialNumber < dockerProperties.getGpuCorePerHost(); serialNumber++) {
                GpuNode gpuNode = new GpuNode(node, serialNumber);
                gpuNodes.put(ip, serialNumber, gpuNode);
            }
        }
    }

    private void initializeNodeUsage() {
        taskScheduler.scheduleWithFixedDelay(() -> {
            syncLock.writeLock().lock();
            try {
                NodeUsage.sync(new ArrayList<>(nodes.values()), taskExecutor);
            } catch (Exception e) {
                log.warn("sync node usage exception", e);
            } finally {
                syncLock.writeLock().unlock();
            }
        }, 300000);
    }

    private void initializeGpuUsageCount() {
        List<Container> gpuContainers = dokkaContainerMapper.listGpuContainers();
        for (Container gpuContainer : gpuContainers) {
            Container container = DockerClient.inspectContainer(gpuContainer.getHost(), gpuContainer.getDockerId());
            if (container != null && "running".equals(container.getStatus())) {
                GpuNode gpuNode = gpuNodes.get(gpuContainer.getHost(), gpuContainer.getGpu());
                if (gpuNode != null) {
                    selectLock.readLock().lock();
                    try {
                        gpuNode.getUsageCount().getAndIncrement();
                    } finally {
                        selectLock.readLock().unlock();
                    }
                }
            }
        }
    }

    @Value("${project.internet}")
    private Boolean recorded;

    private void initializeScheduler() {
        Filter filter = new HealthFilter(recorded);
        Strategy strategy = new ResourcesSpreadStrategy();
        scheduler = new Scheduler(Lists.newArrayList(filter), strategy);
    }

    private void initializeNovncConnector() {
        NovncConnector connector = null;
        if ("local".equals(novncProperties.getMode())) {
            connector = new LocalNovncConnector();
        } else if ("remote".equals(novncProperties.getMode())) {
            connector = new RemoteNovncConnector();
        }
        if (connector == null) {
            log.error("novnc connector mode error");
            System.exit(0);
        }
        connector.initialize();
        this.connector = connector;
    }

    private void initializeDeleteErrorContainer() {
        taskScheduler.scheduleWithFixedDelay(() -> {
            try {
                for (Container errorContainer : dokkaContainerMapper.listErrorContainers()) {
                    DockerClient.stopContainer(errorContainer.getHost(), errorContainer.getDockerId());
                    if (DockerClient.deleteContainer(errorContainer.getHost(), errorContainer.getDockerId())) {
                        dokkaContainerMapper.deleteErrorContainerById(errorContainer.getId());
                    }
                }
            } catch (Exception e) {
                log.error("schedule delete error container exception", e);
            }
        }, 300000);
    }

    private Container createContainer(ContainerConfig containerConfig, Boolean retry) throws DokkaException {
        Node node = selectNode(containerConfig);
        if (node == null) {
            throw new DokkaException("虚拟化服务器集群无可用节点，无法创建实验环境");
        }

        if (!node.tryAcquire()) {
            decrNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), node);
            throw new DokkaException("正在启动实验环境的人数过多，请稍后在进入实验环境");
        }

        String dockerId;
        Result result;
        try {
            dockerId = DockerClient.createContainer(node.getIp(), containerConfig);
            if (StringUtils.isBlank(dockerId)) {
                decrNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), node);
                throw new DokkaException("创建实验环境异常");
            }

            result = DockerClient.startContainer(node.getIp(), dockerId);
            if (result == Result.NO_SUCH_CONTAINER || result == Result.SERVER_ERROR || result == Result.FORBIDDEN) {
                decrNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), node);
                if (!DockerClient.deleteContainer(node.getIp(), dockerId)) {
                    retryDeleteContainer(node.getIp(), dockerId);
                }
            }
        } finally {
            node.releaseAcquire();
        }

        if (result == Result.NO_SUCH_CONTAINER || result == Result.SERVER_ERROR || result == Result.FORBIDDEN) {
            if (retry) {
                return createContainer(containerConfig, false);
            } else {
                throw new DokkaException("创建实验环境异常");
            }
        }

        Container container = DockerClient.inspectContainer(node.getIp(), dockerId);
        if (container == null) {
            decrNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), node);
            throw new DokkaException("创建实验环境异常");
        }

        container.setHost(node.getIp());
        container.setGpu(-1);
        if ("unknown".equals(container.getStatus()) || container.getVncExport() == null) {
            if (container.getVncExport() == null) {
                log.error("get vnc export failed");
            }
            rollBackContainerCreation(containerConfig, node, container);
            throw new DokkaException("创建实验环境异常");
        }

        try {
            dokkaContainerMapper.insert(container);
        } catch (Exception e) {
            log.error("insert container failed", e);
            rollBackContainerCreation(containerConfig, node, container);
            throw new DokkaException("创建实验环境异常");
        }

        String novnc = connector.getNovnc(container.getId(), node.getIp() + ":" + container.getVncExport());
        if (novnc == null) {
            log.error("get novnc failed");
            rollBackContainerCreation(containerConfig, node, container);
            try {
                dokkaContainerMapper.deleteById(container.getId());
            } catch (Exception e) {
                log.error("delete container {} failed", container.getId(), e);
            }
            throw new DokkaException("创建实验环境异常");
        }
        container.setNovnc(novnc);
        return container;
    }

    private static final String JUPYTER_GPU_LIMIT_COMMAND = "echo -n \"#!/usr/bin/env python\n" +
            "# coding: utf-8\n" +
            "\n" +
            "import tensorflow as tf\n" +
            "from keras.backend.tensorflow_backend import set_session\n" +
            "config = tf.ConfigProto()\n" +
            "config.gpu_options.allocator_type = 'BFC'\n" +
            "config.gpu_options.per_process_gpu_memory_fraction = %s\n" +
            "config.gpu_options.allow_growth = True\n" +
            "set_session(tf.Session(config = config))\" > /root/.ipython/profile_default/startup/startup.py";

    private Container createContainerWithGpu(ContainerConfig containerConfig, String gpuPercent) throws DokkaException {
        GpuNode gpuNode = selectGpuNode(containerConfig);
        if (gpuNode == null) {
            throw new DokkaException("虚拟化服务器集群无可用节点，无法创建实验环境");
        }

        if (!gpuNode.getNode().tryAcquire()) {
            decrGpuNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), gpuNode);
            throw new DokkaException("正在启动实验环境的人数过多，请稍后在进入实验环境");
        }

        String dockerId;
        try {
            dockerId = gpuNode.getNode().createContainerWithGpu(containerConfig, gpuNode.getSerialNumber(), true);
            if (StringUtils.isBlank(dockerId)) {
                log.error("create container failed: gpu container return dockerId null");
                decrGpuNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), gpuNode);
                throw new DokkaException("创建实验环境异常");
            }
        } finally {
            gpuNode.getNode().releaseAcquire();
        }

        Container container = DockerClient.inspectContainer(gpuNode.getNode().getIp(), dockerId);
        if (container == null) {
            decrGpuNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), gpuNode);
            throw new DokkaException("创建实验环境异常");
        }

        container.setHost(gpuNode.getNode().getIp());
        container.setGpu(gpuNode.getSerialNumber());
        if (!"running".equals(container.getStatus())) {
            log.error("start container failed");
            rollBackGpuContainerCreation(containerConfig, gpuNode, container);
            throw new DokkaException("创建实验环境异常");
        }
        if (container.getVncExport() == null) {
            log.error("get vnc export failed");
            rollBackGpuContainerCreation(containerConfig, gpuNode, container);
            throw new DokkaException("创建实验环境异常");
        }

        if (StringUtils.isNotBlank(gpuPercent)) {
            String execId = DockerClient.createExec(gpuNode.getNode().getIp(), container.getDockerId(),
                    String.format(JUPYTER_GPU_LIMIT_COMMAND, gpuPercent));
            if (execId == null || !DockerClient.startExec(gpuNode.getNode().getIp(), execId)) {
                rollBackGpuContainerCreation(containerConfig, gpuNode, container);
                throw new DokkaException("创建实验环境异常");
            }
        }

        try {
            dokkaContainerMapper.insert(container);
        } catch (Exception e) {
            log.error("insert container failed", e);
            rollBackGpuContainerCreation(containerConfig, gpuNode, container);
            throw new DokkaException("创建实验环境异常");
        }

        String novnc = connector.getNovnc(container.getId(), gpuNode.getNode().getIp() + ":" + container.getVncExport());
        if (novnc == null) {
            log.error("get novnc failed");
            rollBackGpuContainerCreation(containerConfig, gpuNode, container);
            try {
                dokkaContainerMapper.deleteById(container.getId());
            } catch (Exception e) {
                log.error("delete container {} failed", container.getId(), e);
            }
            throw new DokkaException("创建实验环境异常");
        }
        container.setNovnc(novnc);
        return container;
    }

    private Node selectNode(ContainerConfig containerConfig) {
        selectLock.writeLock().lock();
        try {
            Integer containerQuota = dokkaConfigService.getContainerQuota();
            if (containerQuota != null && getRunningCount() >= containerQuota) {
                throw new DokkaException("运行容器数已达到设置最大数，无法创建更多实验环境");
            }

            Node node = scheduler.selectNodeForContainers(containerConfig, new ArrayList<>(nodes.values()));
            if (node != null) {
                node.getNodeUsage().update(containerConfig.getCpu(), containerConfig.getMemory(), 1);
                return node;
            }
        } finally {
            selectLock.writeLock().unlock();
        }
        return null;
    }

    private GpuNode selectGpuNode(ContainerConfig containerConfig) {
        selectLock.writeLock().lock();
        try {
            Integer containerQuota = dokkaConfigService.getContainerQuota();
            if (containerQuota != null && getRunningCount() >= containerQuota) {
                throw new DokkaException("运行容器数已达到设置最大数，无法创建更多实验环境");
            }

            GpuNode gpuNode = scheduler.selectGpuNodeForContainers(containerConfig, new ArrayList<>(gpuNodes.values()));
            if (gpuNode != null) {
                gpuNode.getNode().getNodeUsage().update(containerConfig.getCpu(), containerConfig.getMemory(), 1);
                gpuNode.getUsageCount().getAndIncrement();
                return gpuNode;
            }
        } finally {
            selectLock.writeLock().unlock();
        }
        return null;
    }

    private void incrNodeUsage(Double cpu, Long memory, Node node) {
        if (node != null && node.getNodeUsage() != null) {
            selectLock.readLock().lock();
            try {
                node.getNodeUsage().update(cpu, memory, 1);
            } finally {
                selectLock.readLock().unlock();
            }
        }
    }

    private void decrNodeUsage(Double cpu, Long memory, Node node) {
        if (node != null && node.getNodeUsage() != null) {
            selectLock.readLock().lock();
            try {
                node.getNodeUsage().update(-1 * cpu, -1 * memory, -1);
            } finally {
                selectLock.readLock().unlock();
            }
        }
    }

    private void incrGpuNodeUsage(Double cpu, Long memory, GpuNode gpuNode) {
        selectLock.readLock().lock();
        try {
            if (gpuNode != null) {
                gpuNode.getUsageCount().getAndIncrement();
                if (gpuNode.getNode() != null && gpuNode.getNode().getNodeUsage() != null) {
                    gpuNode.getNode().getNodeUsage().update(cpu, memory, 1);
                }
            }
        } finally {
            selectLock.readLock().unlock();
        }
    }

    private void decrGpuNodeUsage(Double cpu, Long memory, GpuNode gpuNode) {
        selectLock.readLock().lock();
        try {
            if (gpuNode != null) {
                gpuNode.getUsageCount().getAndDecrement();
                if (gpuNode.getNode() != null && gpuNode.getNode().getNodeUsage() != null) {
                    gpuNode.getNode().getNodeUsage().update(-1 * cpu, -1 * memory, -1);
                }
            }
        } finally {
            selectLock.readLock().unlock();
        }
    }

    private void rollBackContainerCreation(ContainerConfig containerConfig, Node node, Container container) {
        decrNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), node);
        if (DockerClient.stopContainer(container.getHost(), container.getDockerId()) == Result.SERVER_ERROR) {
            incrNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), node);
        }
        if (!DockerClient.deleteContainer(container.getHost(), container.getDockerId())) {
            retryDeleteContainer(container.getHost(), container.getDockerId());
        }
    }

    private void rollBackGpuContainerCreation(ContainerConfig containerConfig, GpuNode gpuNode, Container container) {
        decrGpuNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), gpuNode);
        if (DockerClient.stopContainer(container.getHost(), container.getDockerId()) == Result.SERVER_ERROR) {
            incrGpuNodeUsage(containerConfig.getCpu(), containerConfig.getMemory(), gpuNode);
        }
        if (!DockerClient.deleteContainer(container.getHost(), container.getDockerId())) {
            retryDeleteContainer(container.getHost(), container.getDockerId());
        }
    }

    private void retryDeleteContainer(String host, String dockerId) {
        dokkaContainerMapper.insertErrorContainer(host, dockerId);
    }

    @Override
    public Integer getRunningCount() {
        int count = 0;
        for (Node node : nodes.values()) {
            if (node.getNodeUsage() != null) {
                count += node.getNodeUsage().getRunningCount();
            }
        }
        return count;
    }

    private Container getContainerById(Long dokkaId) {
        Container container;
        try {
            container = dokkaContainerMapper.getById(dokkaId);
        } catch (Exception e) {
            log.error("query container {} failed", dokkaId, e);
            return null;
        }
        if (container == null || !nodes.containsKey(container.getHost())) {
            return null;
        }

        Container inspectContainer = DockerClient.inspectContainer(container.getHost(), container.getDockerId());
        if (inspectContainer == null) {
            try {
                dokkaContainerMapper.deleteById(dokkaId);
            } catch (Exception e) {
                log.error("delete container {} failed", dokkaId, e);
            }
            return null;
        }

        inspectContainer.setId(container.getId());
        inspectContainer.setHost(container.getHost());
        inspectContainer.setGpu(container.getGpu());
        return inspectContainer;
    }

    private void incrNodeUsage(Container container) {
        if (container.getGpu() >= 0) {
            GpuNode gpuNode = gpuNodes.get(container.getHost(), container.getGpu());
            incrGpuNodeUsage(container.getCpu(), container.getMemory(), gpuNode);
        } else {
            Node node = nodes.get(container.getHost());
            incrNodeUsage(container.getCpu(), container.getMemory(), node);
        }
    }

    private void decrNodeUsage(Container container) {
        if (container.getGpu() >= 0) {
            GpuNode gpuNode = gpuNodes.get(container.getHost(), container.getGpu());
            decrGpuNodeUsage(container.getCpu(), container.getMemory(), gpuNode);
        } else {
            Node node = nodes.get(container.getHost());
            decrNodeUsage(container.getCpu(), container.getMemory(), node);
        }
    }
}
